import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import make_interp_spline

# ================================================== 画图 ===============================================
def plot_by_matplot(path_csv):
    # 1. 读取CSV数据
    df = pd.read_csv(path_csv)  # 假设文件名为drone_trajectory.csv

    # 2. 绘制轨迹图
    plt.figure(figsize=(12, 8))
    plt.plot(df['longitude'][:80], df['latitude'][:80], 'o-', color='#1f77b4', markersize=5, linewidth=2)

    # 3. 添加地理参考线（可选）
    plt.axhline(y=37.77, color='gray', linestyle='--', alpha=0.5)  # 参考纬度线
    plt.axvline(x=-122.42, color='gray', linestyle='--', alpha=0.5)  # 参考经度线

    # 4. 设置图形属性
    plt.title('Drone Flight Trajectory', fontsize=14)
    plt.xlabel('Longitude', fontsize=12)
    plt.ylabel('Latitude', fontsize=12)
    plt.grid(True, linestyle=':', alpha=0.7)
    plt.grid(True, linestyle=':', alpha=0.7)

    # 5. 显示/保存图像
    # plt.savefig('flight_trajectory.png', dpi=300, bbox_inches='tight')
    plt.show()


def plot_ori_smooth(df, df_smooth):
    plt.figure(figsize=(15, 10))
    # 原始轨迹（红色虚线）
    plt.plot(df['x'], df['y'], 'ro--', markersize=4, label='origin', alpha=0.6)
    # 平滑轨迹（蓝色实线）
    plt.plot(df_smooth['x'], df_smooth['y'], 'bo-', linewidth=2, label='smooth', alpha=0.9)
    # 边界参考线
    plt.axhline(y=0, color='gray', linestyle=':', alpha=0.5)
    plt.axhline(y=360, color='gray', linestyle=':', alpha=0.5)
    plt.axvline(x=0, color='gray', linestyle=':', alpha=0.5)
    plt.axvline(x=360, color='gray', linestyle=':', alpha=0.5)
    # 图形标注
    plt.title('Drone Flight Path in 360x360 Grid (with smooth)', fontsize=16)
    plt.xlabel('Longitude Scaled (0-360)', fontsize=12)
    plt.ylabel('Latitude Scaled (0-360)', fontsize=12)
    plt.legend()
    plt.grid(True, linestyle=':', alpha=0.7)
    # plt.savefig('smoothed_trajectory.png', dpi=300, bbox_inches='tight')
    plt.show()


################################################ 转换函数 #############################################
# 转换函数:经纬度转(x,y)（线性缩放至360x360）
def latlon_to_xy(lat, lon, lat_min, lat_max, lon_min, lon_max):
    # 经度方向（x轴）：从西到东（最小→最大）
    x = 260 * abs(lon - lon_min) / abs(lon_max - lon_min) + 45
    # 纬度方向（y轴）：从南到北（最小→最大）
    y = 260 * abs(lat - lat_min) / abs(lat_max - lat_min) + 40
    return x, y


# 逆转换函数:(x,y)转经纬度
def xy_to_latlon(x, y, lat_min, lat_max, lon_min, lon_max):
    # 逆向线性缩放公式
    # x = 260 * abs(lon - lon_min) / abs(lon_max - lon_min) + 45
    # -> (x - 45) * abs(lon_max - lon_min) / 260 = abs(lon - lon_min)
    # -> lon = lon_min + (x - 45) * (lon_max - lon_min) / 260
    lon = lon_min + (x - 45) * (lon_max - lon_min) / 260
    lat = lat_min + (y - 40) * (lat_max - lat_min) / 260
    return lat, lon


# 转换函数:(x,y)转(theta,phi)
def xy_to_polar(x, y):
    # 中心平移
    x_mid = x - 180
    y_mid = y - 180
    # 计算半径（防除零处理）
    r = np.sqrt(x_mid ** 2 + y_mid ** 2)
    r = np.where(r == 0, 1e-10, r)  # 当x=180且y=180时，设置极小值
    # 计算th（强制范围约束）
    th_ratio = np.clip(r / 180, -1.0, 1.0)  # 确保arcsin输入有效
    th_rad = np.arcsin(th_ratio)
    th_deg = np.degrees(th_rad)
    # 计算ph（带象限判断）
    ph_rad = np.arccos(x_mid / r)
    ph_deg = np.degrees(ph_rad)
    # 处理y_mid负向情况
    ph_deg = np.where(y_mid < 0, ph_deg + 180, ph_deg)
    # 角度范围修正（0-360度）
    th_deg = np.mod(th_deg, 360)
    ph_deg = np.mod(ph_deg, 360)
    return th_deg, ph_deg


# 逆转换函数:(theta,phi)转(x,y)
def polar_to_xy(th_deg, ph_deg):
    # 将角度转回弧度
    th_rad = np.radians(th_deg)
    ph_rad = np.radians(ph_deg)

    # 计算半径 r
    # th = arcsin(r/180) -> sin(th) = r/180 -> r = 180*sin(th)
    r = 180 * np.sin(th_rad)

    # 计算平移后的 x 和 y 坐标 (x_mid, y_mid)
    x_mid = r * np.cos(ph_rad)
    # 根据 ph 的象限恢复 y_mid 的符号
    # 当 ph 在 [0, 180) 时，y_mid >= 0; 当 ph 在 [180, 360) 时，y_mid < 0
    y_mid_sign = np.where((ph_deg >= 0) & (ph_deg < 180), 1, -1)
    y_mid = y_mid_sign * np.sqrt(np.maximum(r ** 2 - x_mid ** 2, 0))  # 使用 maximum 避免负数开方

    original_arccos_value = np.where(ph_deg > 180, ph_deg - 180, ph_deg)
    original_arccos_rad = np.radians(original_arccos_value)
    x_mid_corrected = r * np.cos(original_arccos_rad)

    # y_mid 的符号由 ph_deg > 180 判断
    y_mid_sign = np.where(ph_deg > 180, -1, 1)  # 如果 ph_deg > 180，则原始 y_mid < 0
    y_mid_corrected = y_mid_sign * np.sqrt(np.maximum(r ** 2 - x_mid_corrected ** 2, 0))

    # 恢复原始的 x 和 y 坐标
    x_corrected = x_mid_corrected + 180
    y_corrected = y_mid_corrected + 180

    return x_corrected, y_corrected


################################################ 主函数 #############################################
# 读取数据集, 转(x,y), 插值, 转(theta,phi)
def latlon_2_xy_interp(path_csv, path_csv_res, path_csv_recover):
    # 1. 读取数据并转换到360x360坐标系
    df = pd.read_csv(path_csv)

    lon_min, lon_max = df['longitude'].min(), df['longitude'].max()
    lat_min, lat_max = df['latitude'].min(), df['latitude'].max()

    df[['x', 'y']] = df.apply(lambda row: latlon_to_xy(row['latitude'], row['longitude'],
                                                         lat_min, lat_max, lon_min, lon_max),
                              axis=1, result_type='expand')
    # 2. 生成平滑插值数据
    # 原始数据索引作为插值参数
    x_original = df['x'].values
    y_original = df['y'].values
    t = np.linspace(0, 1, len(x_original))  # 归一化参数
    # 创建三次样条插值函数
    spline_x = make_interp_spline(t, x_original, k=3)  # 三次样条（k=3）
    spline_y = make_interp_spline(t, y_original, k=3)
    # 生成更密集的插值点（增加5倍采样）
    t_new = np.linspace(0, 1, len(x_original) * 5)
    x_smooth = spline_x(t_new)
    y_smooth = spline_y(t_new)

    # 3. 合并原始与平滑数据
    df_smooth = pd.DataFrame({'x': x_smooth, 'y': y_smooth})

    # 4. 可视化对比
    plot_ori_smooth(df, df_smooth)

    # 5. 坐标系转换(x,y)转为(theta,phi)
    th_smooth, ph_smooth = xy_to_polar(x_smooth, y_smooth)

    # 6. 从(theta,phi)转换回(x,y)和经纬度
    # 俯仰角和方位角转换为直角坐标系(x,y)
    x_recovered, y_recovered = polar_to_xy(th_smooth, ph_smooth)
    # 直角坐标系(x,y)转换回经纬度
    lat_recovered, lon_recovered = xy_to_latlon(x_recovered, y_recovered, lat_min, lat_max, lon_min, lon_max)

    # 7. 创建结果的 DataFrame 并保存
    result_df = pd.DataFrame({
        'x_smooth': x_smooth,
        'y_smooth': y_smooth,
        'th_smooth': th_smooth,  # 离开原点的最大角度（0-90度）
        'ph_smooth': ph_smooth  # 方位角（0-360度）
    })
    result_df.to_csv(path_csv_res, index=False)

    # 8. 创建含恢复结果的 DataFrame 并保存
    recovery_df = pd.DataFrame({
        'recovered_longitude': lon_recovered,
        'recovered_latitude': lat_recovered,
        'recovered_x': x_recovered,
        'recovered_y': y_recovered,
        'th_smooth': th_smooth,
        'ph_smooth': ph_smooth
    })

    recovery_df.to_csv(path_csv_recover, index=False)
    print(f"恢复结果已保存到 {path_csv_recover}")


# 读取latlon_2_xy_interp插值后, (theta,phi)继续插值, 转(x,y), 转经纬度
def read_and_interpolate_csv(path_csv_dataset, path_csv_interp_before, path_csv_interp_after):
    # 1.读取CSV文件
    df = pd.read_csv(path_csv_interp_before)

    # 2.插值
    th_smooth = []
    ph_smooth = []
    # 遍历每一行，计算需要插入的点
    for i in range(len(df) - 1):
        start_th = df.iloc[i]['th_smooth']
        end_th = df.iloc[i + 1]['th_smooth']
        start_ph = df.iloc[i]['ph_smooth']
        end_ph = df.iloc[i + 1]['ph_smooth']
        # 计算需要插入的点的数量
        num_points = int(max(abs(end_th - start_th), abs(end_ph - start_ph)))
        if num_points > 0:
            th_values = np.linspace(start_th, end_th, num_points + 1)
            ph_values = np.linspace(start_ph, end_ph, num_points + 1)

            for j in range(num_points + 1):
                th_smooth.append(th_values[j])
                ph_smooth.append(ph_values[j])
        else:
            th_smooth.append(start_th)
            ph_smooth.append(start_ph)
    # 添加最后一个点
    th_smooth.append(df.iloc[-1]['th_smooth'])
    ph_smooth.append(df.iloc[-1]['ph_smooth'])
    th_smooth = np.array(th_smooth)
    ph_smooth = np.array(ph_smooth)

    # 3. 从(theta,phi)转换回(x,y)和经纬度
    df = pd.read_csv(path_csv_dataset)
    lon_min, lon_max = df['longitude'].min(), df['longitude'].max()
    lat_min, lat_max = df['latitude'].min(), df['latitude'].max()
    # 俯仰角和方位角转换为直角坐标系(x,y)
    x_smooth, y_smooth = polar_to_xy(th_smooth, ph_smooth)
    # 直角坐标系(x,y)转换回经纬度
    lat_smooth, lon_smooth = xy_to_latlon(x_smooth, y_smooth, lat_min, lat_max, lon_min, lon_max)

    # 4. 创建含恢复结果的 DataFrame 并保存
    recovery_df = pd.DataFrame({
        'lon_smooth': lon_smooth,
        'lat_smooth': lat_smooth,
        'x_smooth': x_smooth,
        'y_smooth': y_smooth,
        'th_smooth': th_smooth,
        'ph_smooth': ph_smooth
    })

    recovery_df.to_csv(path_csv_interp_after, index=False)
    print(f"恢复结果已保存到 {path_csv_interp_after}")





if __name__ == "__main__":
    latlon_2_xy_interp("../../../../files/archive/uav_navigation_dataset.csv",
                       "../../../files/archive/uav_dataset_xy_thph_3.csv",
                       "../../../files/archive/uav_dataset_recovery.csv")
    read_and_interpolate_csv("../../../../files/archive/uav_navigation_dataset.csv",
                       "../../../files/archive/uav_dataset_xy_thph.csv",
                       "../../../files/archive/uav_dataset_thph_xy_lonlat.csv")



